#ifndef	rmsnorm_included
#define	rmsnorm_included

__global__	void	_rmsnormf(uint32_t	r,	half	*inp,	float	*norm) {
    uint32_t	id=blockIdx.x*blockDim.x+threadIdx.x;
    float	nor=0;
    half	*in=inp+id*r;
    for(uint32_t	i=0;	i<r;	i+=8) {
        float8	z=_b2f8(*(half8*)(in+i));
        nor+=_fdot8(z,z);
    }
    norm[id]=nor;
    nor=sqrtf(r/nor);
    for(uint32_t	i=0;	i<r;	i+=8)	*(half8*)(in+i)=_f2b8(_fmul18(_b2f8(*(half8*)(in+i)),nor));
}

__global__	void	_rmsnormb(uint32_t	r,	half	*inp,	half	*gin,	float	*norm) {
    uint32_t	id=blockIdx.x*blockDim.x+threadIdx.x;
    float	mg=0,	s=sqrtf(r/norm[id]);
    half	*in=inp+id*r,	*gi=gin+id*r;
    for(uint32_t	i=0;	i<r;	i+=8) mg+=_fdot8(_b2f8(*(half8*)(gi+i)),_b2f8(*(half8*)(in+i)));
    mg/=norm[id]*s;
    for(uint32_t	i=0;	i<r;	i+=8)	*(half8*)(gi+i)=_f2b8(_fadd8(_fmul18(_b2f8(*(half8*)(gi+i)),s),_fmul18(_b2f8(*(half8*)(in+i)),-mg)));
}

template<uint32_t	R,	uint32_t	C,	uint32_t	H>
struct	rmsnorm {
    tensor<float,C*H>	nor;

    void	fw(tensor<half,R*C>	&inp) {
        const	uint32_t	warp=warp_size(1,C*H);
        _rmsnormf<<<C*H/warp,warp>>>(R/H,inp.p,nor.p);
    }

    void	bk(tensor<half,R*C>	&inp,	tensor<half,R*C>	&gin) {
        const	uint32_t	warp=warp_size(1,C*H);
        _rmsnormb<<<C*H/warp,warp>>>(R/H,inp.p,gin.p,nor.p);
    }
};

#endif
